import Foundation
import CommonCrypto

// 扩展Data处理十六进制字符串
extension Data {
    init?(hexString: String) {
        let length = hexString.count / 2
        var data = Data(capacity: length)
        for i in 0..<length {
            let start = hexString.index(hexString.startIndex, offsetBy: i*2)
            let end = hexString.index(start, offsetBy: 2)
            let bytes = hexString[start..<end]
            if let num = UInt8(bytes, radix: 16) {
                data.append(num)
            } else {
                return nil
            }
        }
        self = data
    }
}

// PKCS7 填充
func pkcs7Padding(data: Data, blockSize: Int) -> Data {
    let paddingLength = blockSize - (data.count % blockSize)
    let padding = Data(repeating: UInt8(paddingLength), count: paddingLength)
    return data + padding
}

// 移除 PKCS7 填充
func removePkcs7Padding(data: Data) -> Data? {
    guard let lastByte = data.last else { return nil }
    let paddingLength = Int(lastByte)
    let unpaddedData = data.dropLast(paddingLength)
    return unpaddedData
}

// AES加密函数
func encryptAES_CFB(plainText: String, key: Data, iv: Data) -> String? {
    let plainData = pkcs7Padding(data: Data(plainText.utf8), blockSize: kCCBlockSizeAES128)
    
    var cryptor: CCCryptorRef?
    let status = key.withUnsafeBytes { keyBytes in
        iv.withUnsafeBytes { ivBytes in
            CCCryptorCreateWithMode(
                CCOperation(kCCEncrypt),
                CCMode(kCCModeCFB),
                CCAlgorithm(kCCAlgorithmAES),
                CCPadding(ccNoPadding), // CFB模式不需要填充
                ivBytes.baseAddress,
                keyBytes.baseAddress,
                key.count,
                nil,
                0,
                0,
                CCModeOptions(kCCModeOptionCTR_BE),
                &cryptor)
        }
    }
    
    guard status == kCCSuccess, let cryptor = cryptor else { return nil }
    
    defer { CCCryptorRelease(cryptor) }
    
    let bufferSize = CCCryptorGetOutputLength(cryptor, plainData.count, true)
    var outputData = Data(count: bufferSize)
    var outputLength = 0
    
    let updateStatus = plainData.withUnsafeBytes { plainBytes in
        outputData.withUnsafeMutableBytes { outputBytes in
            CCCryptorUpdate(
                cryptor,
                plainBytes.baseAddress,
                plainData.count,
                outputBytes.baseAddress,
                bufferSize,
                &outputLength
            )
        }
    }
    
    guard updateStatus == kCCSuccess else { return nil }
    
    outputData.removeSubrange(outputLength..<outputData.count)
    return outputData.base64EncodedString()
}

// AES解密函数
func decryptAES_CFB(cipherText: String, key: Data, iv: Data) -> String? {
    guard let cipherData = Data(base64Encoded: cipherText) else { return nil }
    
    var cryptor: CCCryptorRef?
    let status = key.withUnsafeBytes { keyBytes in
        iv.withUnsafeBytes { ivBytes in
            CCCryptorCreateWithMode(
                CCOperation(kCCDecrypt),
                CCMode(kCCModeCFB),
                CCAlgorithm(kCCAlgorithmAES),
                CCPadding(ccNoPadding), // CFB模式不需要填充
                ivBytes.baseAddress,
                keyBytes.baseAddress,
                key.count,
                nil,
                0,
                0,
                CCModeOptions(kCCModeOptionCTR_BE),
                &cryptor)
        }
    }
    
    guard status == kCCSuccess, let cryptor = cryptor else { return nil }
    
    defer { CCCryptorRelease(cryptor) }
    
    let bufferSize = CCCryptorGetOutputLength(cryptor, cipherData.count, true)
    var outputData = Data(count: bufferSize)
    var outputLength = 0
    
    let updateStatus = cipherData.withUnsafeBytes { cipherBytes in
        outputData.withUnsafeMutableBytes { outputBytes in
            CCCryptorUpdate(
                cryptor,
                cipherBytes.baseAddress,
                cipherData.count,
                outputBytes.baseAddress,
                bufferSize,
                &outputLength
            )
        }
    }
    
    guard updateStatus == kCCSuccess else { return nil }
    
    outputData.removeSubrange(outputLength..<outputData.count)
    return removePkcs7Padding(data: outputData).flatMap { String(data: $0, encoding: .utf8) }
}

// 测试代码 neAxw1KQ02SHJQquPe7OCQ==
let aesKeyHex = "3f0ad7320645e75a8fa310e5aa7006bca056869b5e0be1fd938f5248ba6a17a6"
let ivString = "aybolfhuangWeb3%"

guard let aesKey = Data(hexString: aesKeyHex),
      let ivData = ivString.data(using: .utf8) else {
    fatalError("Invalid key or IV")
}

let plainText = "asdsadsadsadsa"

// 加密
if let encrypted = encryptAES_CFB(plainText: plainText, key: aesKey, iv: ivData) {
    print("Encrypted (Base64):", encrypted)
    
    // 解密
    if let decrypted = decryptAES_CFB(cipherText: encrypted, key: aesKey, iv: ivData) {
        print("Decrypted:", decrypted)
    } else {
        print("Decryption failed")
    }
} else {
    print("Encryption failed")
}
